package gqz.avdemo.pusher;

import android.hardware.display.DisplayManager;
import android.hardware.display.VirtualDisplay;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.projection.MediaProjection;
import android.util.Log;
import android.view.Surface;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Arrays;

/**
 * Description
 *
 * @author gqz
 * @date 2020/12/21
 */
public class Codec extends Thread {
   private final String TAG = "Codec-->";

   SocketLive live;
   MediaProjection projection;
   MediaCodec mediaCodec;

   int width = 1080, height = 1920, frameRate = 15, i_rate = 5;
   private VirtualDisplay virtualDisplay;

   public Codec(SocketLive live, MediaProjection projection) {
      this.live = live;
      this.projection = projection;
   }

   public void startPush() {
      try {
         mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
         MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width
                 , height);
         format.setInteger(MediaFormat.KEY_COLOR_FORMAT,
                 MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
         format.setInteger(MediaFormat.KEY_BIT_RATE, width * height / 10);
         format.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
         format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, i_rate);
         mediaCodec.configure(format,null,null,MediaCodec.CONFIGURE_FLAG_ENCODE);
         Surface surface = mediaCodec.createInputSurface();
         virtualDisplay = projection.createVirtualDisplay("*vDisplay", width,
                 height, 1,
                 DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC, surface, null, null);
      } catch (IOException e) {
         e.printStackTrace();
      }

      start();
   }

   @Override
   public void run() {
      super.run();
      mediaCodec.start();
      MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
      while (true) {

         try {
            int bufIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, 100000);
            if (bufIndex >= 0) {
               ByteBuffer buf = mediaCodec.getOutputBuffer(bufIndex);
               dealFrame(buf, bufferInfo);
               mediaCodec.releaseOutputBuffer(bufIndex, false);
            }
         } catch (Exception e) {
            e.printStackTrace();
         }

      }
   }

   public static final int NAL_I = 19;
   public static final int NAL_VPS = 32;//hevc
   public static final int NAL_SPS = 32;//avc
   private byte[] configBuf;

   private void dealFrame(ByteBuffer buf, MediaCodec.BufferInfo bufferInfo) {
      int offset = 4;
      if (buf.get(2) == 0x01) {//判断分隔符
         offset = 3;
      }
      int type = (buf.get(offset) & 0x7e) >> 1;//判断帧类型
      if (type == NAL_VPS) {//保存config帧
         configBuf = new byte[bufferInfo.size];
         buf.get(configBuf);
      } else if (type == NAL_I) {//将config帧拼到I帧前发送
         byte[] data = new byte[bufferInfo.size];
         buf.get(data);
         byte[] newData = new byte[configBuf.length + bufferInfo.size];
         System.arraycopy(configBuf, 0, newData, 0, configBuf.length);
         System.arraycopy(data, 0, newData, configBuf.length, data.length);
         live.sendData(newData);
         Log.v(TAG, "视频帧数据 ->  I ");
      } else {//BP帧直接发
         byte[] data = new byte[bufferInfo.size];
         buf.get(data);
         live.sendData(data);
         Log.v(TAG, "视频帧数据 -> " + Arrays.toString(data));
      }
   }

   private long generatePTS(int index) {
      return 132 + index * 1000000 / frameRate;
   }

   public void onDestroy() {
      if (mediaCodec != null) {
         mediaCodec.release();
         mediaCodec.stop();
      }
   }
}
